#####################################################
# Decision Making in Intelligent Systems            #
# Assignment 2 - Revised Mountain Car Task          #
#                                                   #
# Chaim Bastiaan [5742889] & Kai Krabben [5743036]  #
#####################################################

import math
from random import random
from numpy import *
from copy import deepcopy
from pylab import plot, show, imshow, figure, cm

# action space: 
# +1 is full throttle forward
# -1 is full throttle reverse
# 0 is zero throttle
A = (1, 0, -1)

# set bounds for x and v
xMin = -1.2
xMax = 0.5
vMin = -0.07
vMax = 0.07

# state space: s = (v,t)
# represented by position x and velocity v
    
def terminal(s):
    # state is terminal if right bound has been reached
    #return s[1]==xMax or s[1]==xMin
    return s[1]>=xMax

# transition model
def transitionModel((v0,x0), a):
    x1 = x0 + v0
    v1 = v0 + 0.001*a - 0.0028*math.cos(2*x0 - 0.5)

    # bound operation
    v1 = max(v1, vMin)
    v1 = min(v1, vMax)
    
    if x1 <= xMin:
        # left bound reached
        v1 = 0
        #x1 = xMin
        x1 = xMin+.05
    if x1 >= xMax:
        # right bound reached: goal reached, episode terminated
        pass
    
    x1 = min(x1, xMax)
    #return (x1,v1)
    return (v1,x1)

# reward function
def reward(s0, s1, a):
    return -1

class Grid:
    # represents one discretisation of the state space in 2D grid
    def __init__(self, tilesx, tilesy):
        self.tx = tilesx
        self.ty = tilesy
        self.tw = (vMax-vMin)/float(self.tx)
        self.th = (xMax-xMin)/float(self.ty)
        self.hasoffset = False
        cellList = [0]*(tilesx*tilesy)

    def setOffset(self, valx, valy):
        self.hasoffset = True
        self.offsetx = valx
        self.offsety = valy

    def setCell(self, c, val):
        cellList[self.cellToIndex(c)]=val

    def getCell(self, s):
        if self.hasoffset:
            #c = (math.floor(((s[0]+self.offsetx)/self.tw)), math.floor(((s[1]+self.offsety)/self.th)))
            c = (math.floor((s[0]-vMin)/self.tw), math.floor((s[1]-xMin)/self.th))

        else:
            c = (math.floor((s[0]-vMin)/self.tw), math.floor((s[1]-xMin)/self.th))
        
        return c

    def cellToIndex(self, (x,y)):
        return (x)*self.tx+y

    def getCells(self):
        # return all cells
        for x in xrange(0,self.tx):
          for y in xrange(0,self.ty):
            yield (float(x),float(y))

    def getFeatures(self, s):
        cells = [c==self.getCell(s) for c in self.getCells()]
        return cells



#####################################################
# Learning Algorithm                                #
#####################################################

def makeTiling(grid, N):
    # make tiling of N copies of grid with different random offsets
    tiling = [grid]
    for i in xrange(N):
        nextGrid = Grid(grid.tx,grid.ty)
        nextGrid.setOffset(random.uniform(-grid.tw, grid.tw), random.uniform(-grid.th, grid.th))
        tiling.append(nextGrid)
    return tiling

def makeTheta(tx, ty, N):
    theta = dict()
    theta_a = [0] * tx * ty* (N+1)
    for a in A:
        theta[a] = theta_a
    return theta

def merge(seq):
    merged = []
    for s in seq:
        for x in s:
            merged.append(x)
    return merged

def getFeatures(s, tiling):
    Fa = []
    for grid in tiling:
        Fa.append(grid.getFeatures(s))
    return merge(Fa)

def summedValues(values, indices):
    summedValues = 0
    for i in xrange(len(indices)):
        if indices[i]:
            summedValues += values[i]
    return summedValues

def SarsaL(lam): # lam = lambda
    # Parameters (change if you like)
    tx = 10
    ty = 10
    ntilings = 1 # number of tilings, 0 for 'discretization' by grid, only
    episodes = 10
    epsilon = 0.2
    gamma = 0.9

    try:
        alpha = 0.2/float(ntilings)
    except:
        alpha = .2

    # initialisation
    grid = Grid(tx,ty)
    tiling = makeTiling(grid, ntilings)

    theta = makeTheta(tx, ty, ntilings)

    for epi in xrange(episodes):
        print "episode %d" % epi
        s = (random.uniform(vMin,vMax), random.uniform(xMin,xMax))
        a = A[random.randint(0,len(A))]
        Fa = getFeatures(s, tiling)

        tempcount = 0
        while not terminal(s):
            tempcount += 1
            s1 = transitionModel(s, a)      # Apply transition model
            delta = reward(s, s1, a) - summedValues(theta[a], Fa)

            if random.random() <= (1 - epsilon):
                Qlist = []
                for a in A:
                    Fa = getFeatures(s1,tiling)
                    Qlist.append(summedValues(theta[a], Fa))
                Qa = max(Qlist)
                a = A[Qlist.index(Qa)]
                
            else:
                a = A[random.randint(0, len(A))]
                Fa = getFeatures(s1,tiling)
                Qa = summedValues(theta[a], Fa)

            delta = delta + gamma * Qa
            for i in xrange(len(Fa)):
                if Fa[i]:
                    theta[a][i] += alpha * delta
            
            s = s1

    X = [[0]*10]*10
    Y = [[0]*10]*10
    Z = [[0]*10]*10
    for tx in xrange(grid.tx):
        for ty in xrange(grid.ty):
            id = grid.cellToIndex((tx,ty))
            X[tx][ty] = theta[-1][id]
            #Y[tx][ty] = theta[0][id]
            #Z[tx][ty] = theta[1][id]

    figure()
    imshow(X,interpolation=None, cmap=cm.gray )
   #figure()
   #imshow(Y,interpolation=None, cmap=cm.gray )
   #figure()
   #imshow(Z,interpolation=None, cmap=cm.gray )

    show()
      #for tx in xrange(len(tiling[0].tilesx)):
      #    for ty in xrange(len(tiling[0].tilesy)):
      #        X[tx][ty] = 
      #imshow(X)
            
    return theta, tiling


def drive(s,a, theta,tiling):
    x = 0
    while not terminal(s):
        x+=1
        Qlist = []
        for a in A:
            Fa = getFeatures(s,tiling)
            Qlist.append(summedValues(theta[a], Fa))
        a = A[Qlist.index(max(Qlist))]
        s = transitionModel(s,a)
        #print s,a
        if s[0]<0:
            plot(x,abs(s[1]-xMax),'b.',markersize=(-s[0])*120)
        else:
            plot(x,abs(s[1]-xMax),'r.',markersize=s[0]*120)

    print "Final state and action:"
    print s,a
    show()

print "Performing Sarsa-Lambda"
th,tiling = SarsaL(.2)

tv = 0
tx = (xMin+xMax)/2.0
ta = 1
print "Picking correct actions starting at v:%d, x:%d, a:%d" % (tv, tx, ta)
print "vMin = %f, vMax = %f, xMin = %f, xMax = %f" % (vMin, vMax, xMin, xMax)
drive((tv,tx),ta,th,tiling)
